#include "Avalon/Codecs/AesDecoder.hpp"
#include <cryptopp/dh.h>
#include <cryptopp/hex.h>
#include <cryptopp/sha.h>
#include <cryptopp/modes.h>
#include <cryptopp/osrng.h>
#include "Avalon/Codecs/AesEncoder.hpp"
#include "Avalon/Codecs/CodedReader.hpp"
#include "Avalon/Codecs/CodedWriter.hpp"
#include "Avalon/IO/Buffer.hpp"
#include "Avalon/IO/Channel.hpp"
#include "Avalon/IO/IOException.hpp"

using namespace Avalon;
using namespace Avalon::Codecs;
using namespace Avalon::IO;
using namespace boost;
using namespace boost::posix_time;
using namespace CryptoPP;
using namespace std;

namespace {
  byte* CopyIV(const byte* source) {
    byte* copy = new byte[SHA::DIGESTSIZE];
    memcpy(copy, source, SHA::DIGESTSIZE);
    return copy;
  }
}

Channel* AesDecoder::EstablishClientChannel(Channel* channel,
    bool encodeInPlace) {
  Buffer message;
  channel->GetReader().Read(&message)->Wait(pos_infin);
  CryptoPP::Integer prime(reinterpret_cast<const byte*>(message.GetData()),
    message.GetSize());
  channel->GetReader().Read(&message)->Wait(pos_infin);
  CryptoPP::Integer generator(reinterpret_cast<const byte*>(message.GetData()),
    message.GetSize());
  DH keyExchanger(prime, generator);
  SecByteBlock privateKey(keyExchanger.PrivateKeyLength());
  SecByteBlock publicKey(keyExchanger.PublicKeyLength());
  SecByteBlock secretKey(keyExchanger.AgreedValueLength());
  AutoSeededRandomPool randomNumberGenerator;
  keyExchanger.GenerateKeyPair(randomNumberGenerator, privateKey, publicKey);
  channel->GetWriter().Write(publicKey.BytePtr(),
    static_cast<int>(publicKey.SizeInBytes()))->Wait(pos_infin);
  channel->GetReader().Read(&message)->Wait(pos_infin);
  SecByteBlock serverPublicKey(reinterpret_cast<const byte*>(
    message.GetData()), message.GetSize());
  if(!keyExchanger.Agree(secretKey, privateKey, serverPublicKey)) {
    BOOST_THROW_EXCEPTION(IOException(
      "Server-client encryption keys are not compatible."));
  }

  // Diffie-Hellman key exchange succeeded.  Now generate the AES key.
  SecByteBlock aesKey(SHA256::DIGESTSIZE);
  SHA256().CalculateDigest(aesKey, secretKey, secretKey.size());
  if(channel->GetReader().Read(&message)->Get(pos_infin) != AES::BLOCKSIZE) {
    BOOST_THROW_EXCEPTION(IOException("Invalid AES key."));
  }
  byte iv[AES::BLOCKSIZE];
  memcpy(iv, message.GetData(), AES::BLOCKSIZE);
  AesDecoder* decoder = new AesDecoder(aesKey, iv);
  AesEncoder* encoder = new AesEncoder(aesKey, iv);
  CodedReader* aesReader = new CodedReader(channel->GetReader(), decoder);
  CodedWriter* aesWriter = new CodedWriter(channel->GetWriter(), encoder,
    encodeInPlace);
  Channel* secureChannel = new BasicChannel(channel->GetConnection(), aesReader,
    aesWriter);
  return secureChannel;
}

Channel* AesDecoder::EstablishServerChannel(Channel* channel,
    bool encodeInPlace) {

  // Generate prime number and a generator to use for the Diffie-Hellman
  // key exchange as well as a private, public, and secret key.
  // Then send to the client the prime, the generator and the public key.
  AutoSeededRandomPool randomNumberGenerator;
  DH keyExchanger(static_cast<RandomNumberGenerator&>(randomNumberGenerator),
    128);
  CryptoPP::Integer prime = keyExchanger.GetGroupParameters().GetModulus();
  CryptoPP::Integer generator = keyExchanger.GetGroupParameters().GetSubgroupGenerator();
  SecByteBlock privateKey(keyExchanger.PrivateKeyLength());
  SecByteBlock publicKey(keyExchanger.PublicKeyLength());
  SecByteBlock secretKey(keyExchanger.AgreedValueLength());
  keyExchanger.GenerateKeyPair(randomNumberGenerator, privateKey, publicKey);
  scoped_array<byte> encodedInteger(new byte[prime.ByteCount()]);
  prime.Encode(encodedInteger.get(), prime.ByteCount());
  channel->GetWriter().Write(encodedInteger.get(), prime.ByteCount())->Wait(
    pos_infin);
  encodedInteger.reset(new byte[generator.ByteCount()]);
  generator.Encode(encodedInteger.get(), generator.ByteCount());
  channel->GetWriter().Write(encodedInteger.get(), generator.ByteCount())->Wait(
    pos_infin);
  Buffer message;
  channel->GetReader().Read(&message)->Wait(pos_infin);
  SecByteBlock clientPublicKey(reinterpret_cast<const byte*>(message.GetData()),
    message.GetSize());
  if(!keyExchanger.Agree(secretKey, privateKey, clientPublicKey)) {
    BOOST_THROW_EXCEPTION(IOException(
      "Server-client encryption keys are not compatible."));
  }
  channel->GetWriter().Write(publicKey.BytePtr(),
    static_cast<int>(publicKey.SizeInBytes()))->Wait(pos_infin);

  // Diffie-Hellman key exchange succeeded.  Now generate the AES key.
  SecByteBlock aesKey(SHA256::DIGESTSIZE);
  SHA256().CalculateDigest(aesKey, secretKey, secretKey.size());
  byte iv[AES::BLOCKSIZE];
  randomNumberGenerator.GenerateBlock(iv, AES::BLOCKSIZE);
  channel->GetWriter().Write(iv, AES::BLOCKSIZE)->Wait(pos_infin);
  AesDecoder* decoder = new AesDecoder(aesKey, iv);
  AesEncoder* encoder = new AesEncoder(aesKey, iv);
  CodedReader* aesReader = new CodedReader(channel->GetReader(), decoder);
  CodedWriter* aesWriter = new CodedWriter(channel->GetWriter(), encoder,
    encodeInPlace);
  Channel* secureChannel = new BasicChannel(channel->GetConnection(), aesReader,
    aesWriter);
  return secureChannel;
}

AesDecoder::AesDecoder(const SecByteBlock& key, const byte* iv)
    : m_key(key),
      m_iv(CopyIV(iv)),
      m_decoder(m_key, SHA256::DIGESTSIZE, m_iv.get()) {}

AesDecoder::~AesDecoder() {}

bool AesDecoder::IsDecodeInPlaceSupported() {
  return true;
}

void AesDecoder::Decode(const void* source, int sourceSize, void* destination,
    int destinationSize, int* decodedSize) {
  assert(destinationSize >= sourceSize);
  m_decoder.ProcessData(static_cast<byte*>(destination),
    static_cast<const byte*>(source), sourceSize);
  if(decodedSize != NULL) {
    *decodedSize = sourceSize;
  }
}

void AesDecoder::Decode(const void* source, int sourceSize, Buffer* destination,
    int* decodedSize) {
  int startingPoint = destination->GetSize();
  destination->Grow(sourceSize);
  m_decoder.ProcessData(reinterpret_cast<byte*>(destination->GetMutableData()
    + startingPoint), static_cast<const byte*>(source), sourceSize);
  if(decodedSize != NULL) {
    *decodedSize = sourceSize;
  }
}
